#!/usr/bin/env python

import rospy
import tf2_ros
from geometry_msgs.msg import Twist,PoseStamped, Transform,PoseArray,Pose,PoseWithCovarianceStamped
import numpy as np
from ros_numpy import numpify,msgify
from nav_msgs.msg import Path
from sensor_msgs.msg import PointCloud2
from tf.transformations import quaternion_from_euler,euler_matrix
from sensor_msgs.msg import JointState
from trajectory_msgs.msg import MultiDOFJointTrajectory,MultiDOFJointTrajectoryPoint
import torch
import torch.optim as optim
import matplotlib.pyplot as plt
import time

def euler_position_to_T(euler_zyx,pos=None):
    
    n=euler_zyx.shape[1]
    
    if pos==None:
        p=torch.zeros([3,n])
    else:
        p=pos
    
    angle_z=euler_zyx[0,:]
    angle_y=euler_zyx[1,:]
    angle_x=euler_zyx[2,:]
    
    cos_x, sin_x = torch.cos(angle_x), torch.sin(angle_x)
    cos_y, sin_y = torch.cos(angle_y), torch.sin(angle_y)
    cos_z, sin_z = torch.cos(angle_z), torch.sin(angle_z)

    # Compute the elements of the rotation matrix
    R00 = cos_z * cos_y
    R01 = cos_z * sin_y * sin_x - sin_z * cos_x
    R02 = cos_z * sin_y * cos_x + sin_z * sin_x
    R10 = sin_z * cos_y
    R11 = sin_z * sin_y * sin_x + cos_z * cos_x
    R12 = sin_z * sin_y * cos_x - cos_z * sin_x
    R20 = -sin_y
    R21 = cos_y * sin_x
    R22 = cos_y * cos_x
    T_0 = torch.zeros(n)
    T_1 = torch.ones(n)

    # Construct the rotation matrix R from the elements
    transformationT = torch.stack([torch.stack([R00, R01, R02,p[0,:]]),
                                   torch.stack([R10, R11, R12,p[1,:]]),
                                   torch.stack([R20, R21, R22,p[2,:]]),
                                   torch.stack([T_0, T_0, T_0,T_1])])
    
    T=torch.transpose(transformationT,0,2)
    T=torch.transpose(T,1,2)
    
    return T

def pc_W_to_C(pc_W,T_W_C):
    n_horizon=T_W_C.shape[0]
    pc_size=pc_W.shape[1]
    pc_W_aug=torch.vstack([pc_W,torch.ones(pc_size)])
    pc_W_aug=pc_W_aug.repeat(n_horizon,1,1).float()
    
    pc_C_aug=T_W_C@pc_W_aug
    pc_C=pc_C_aug[:,:3,:]
    
    return pc_C


def calculate_visibility(a,fov,bearing):
    a_vis=torch.tensor(a)
    #fov=torch.pi
    v=1/(1+torch.exp(-a_vis*torch.cos(bearing*torch.pi/fov)))#*(1+torch.exp(-torch.tensor(a_vis)))
    v_max=1/(1+torch.exp(-a_vis))
    v_min=1/(1+torch.exp(a_vis))
    
    v_normalized=(v-v_min)/(v_max-v_min)
    return v_normalized

def vectors_to_euler_zy(v):
    v_norm=torch.sqrt(torch.sum(v*v,axis=0))
    v_normalized=v/(v_norm.repeat(3,1))
    
    r11=v_normalized[0,:]
    r21=v_normalized[1,:]
    r31=v_normalized[2,:]
    
    angle_z=torch.atan2(r21,r11)
    angle_y=torch.atan2(-r31,torch.sqrt(r11**2+r21**2+1e-6))
    
    euler_zy=torch.vstack([angle_z,angle_y])
    
    return euler_zy

def step(A,B,x0,u,p,pc_W,n_horizon):
    
    pc_size=pc_W.shape[1]
    
    state_size=x0.shape[0]
    input_size=u.shape[0]
    
    x_flatten=A@x0+B@u
    
    #process agent 
    x=x_flatten.reshape([n_horizon,state_size]).T
    agent_euler_zyx=torch.zeros([3,n_horizon])
    agent_euler_zyx[:2,:]=x
    
    agent_pos=p

    T_C_W=euler_position_to_T(euler_zyx=agent_euler_zyx,pos=agent_pos)
    T_W_C=torch.inverse(T_C_W)
    

    #process pc
    pc_C=pc_W_to_C(pc_W,T_W_C)
    pc_C_center=torch.mean(pc_C,axis=2)
    
    
    pc_norm=torch.sqrt(torch.sum(pc_C**2,axis=1)).reshape([n_horizon,1,pc_size])
    pc_normalized=pc_C/pc_norm
    cos_bear=pc_normalized[:,0,:]
    bear=torch.acos(cos_bear)
    
    #calculate loss
    
    
    vis_pc=calculate_visibility(a=torch.tensor(1.0),fov=torch.pi*1,bearing=bear)
    
    vis_ob0=calculate_visibility(a=torch.tensor(10.0),fov=torch.pi,bearing=bear[:,0])
    vis_ob=vis_ob0*calculate_visibility(a=torch.tensor(10.0),fov=torch.pi/2,bearing=bear[:,0])
    
    v_flatten=u
    v=v_flatten.reshape([n_horizon,2]).T
    a=v[:,1:]-v[:,:-1]
    
    
    loss=-torch.mean(vis_pc)*1.0-torch.mean(vis_ob)*1.0#+torch.mean(a*a)*1e4+torch.mean(v*v)
    
    return loss,vis_pc,vis_ob


def pc2_msg_to_np(pc_msg):
    pc_np=numpify(pc_msg)
    pc_np= np.stack([pc_np[f] for f in ('x', 'y', 'z')])

    return pc_np


def extend_AB(A,B,n_horizon,dt):
    A_ext=A
    for i in range(1,n_horizon):
        A_ext=torch.vstack([A_ext,torch.matrix_power(A,i+1)])

    B_ext=torch.zeros([2*n_horizon,2*n_horizon])  
    for r in range(n_horizon):
        for c in range(n_horizon):
            if r>=c:
                B_ext[2*r:2*r+2,2*c:2*c+2]=torch.matrix_power(A,r-c)@B
    B_ext=B_ext*dt
    
    return A_ext,B_ext
    

class ViewPlanningNode:
    def __init__(self,n_horizon=10,dt=0.1,state_size=2,input_size=2):
        rospy.init_node('view_planning_node')

        
        A=torch.eye(2)
        B=torch.eye(2)
        A_ext,B_ext=extend_AB(A,B,n_horizon,dt)
        
        self.A_ext=A_ext
        self.B_ext=B_ext
        self.p=torch.zeros([4,n_horizon])

        
        self.x0=torch.zeros([2,1])
        self.u=torch.zeros([2,n_horizon])
        self.n_horizon=n_horizon
        self.pc_W=torch.ones([3,100])
        self.ob_W=torch.ones([3,1])
        self.pc_cnt_list=[]
        #self.update_lm()
        
        rospy.Subscriber('/slot/pointcloud2_visible',PointCloud2,self.pc_cb)
        rospy.Subscriber('/slot/pose_array',PoseArray,self.pos_arr_cb)
        rospy.Subscriber('/slot/target/path',Path,self.target_path_cb)
    
        self.test_pub=rospy.Publisher('/slot/test/pc',PointCloud2,queue_size=10)
        self.view_array_pub=rospy.Publisher('/slot/view_array',PoseArray,queue_size=10)
    
    def update_lm(self):
        
        self.lm_W=torch.hstack([self.ob_W,self.pc_W])
        self.lm_W=self.lm_W.float()
    
    def target_path_cb(self,path_msg):
        target_pose_msg=path_msg.poses[-1].pose
        target_pos=numpify(target_pose_msg)[:3,3]
        self.ob_W=torch.tensor(target_pos).reshape([3,1])
        self.update_lm()
    
    
    def init_u(self,view_center='pc_center'):
        n_horizon=self.n_horizon
        T_C_W=euler_position_to_T(euler_zyx=torch.zeros([3,n_horizon]),pos=self.p)
        T_W_C=torch.inverse(T_C_W)
 
        
        lm_C=pc_W_to_C(self.lm_W,T_W_C)
        lm_C_center=torch.mean(lm_C,axis=2).T
        ob_C=lm_C[:,:,0].T
        
        x_init=vectors_to_euler_zy(v=ob_C)
        
        
        if view_center=='pc_center':
            x_init=vectors_to_euler_zy(v=lm_C_center)
        elif view_center=='ob_center':
            x_init=vectors_to_euler_zy(v=ob_C+lm_C_center*0.1)
        else:
            x_init=vectors_to_euler_zy(v=lm_C_center)*0.0
        
        
        x_init_flatten=x_init.T.reshape([2*n_horizon,1])
        
        u_init_flatten=torch.pinverse(self.B_ext)@(x_init_flatten-self.A_ext@self.x0)
        
        u_init_noise=torch.rand(u_init_flatten.shape)*0.
        u_init_flatten+=u_init_noise
        
        u_init_flatten.requires_grad=True

        
        return u_init_flatten
        
        
    def pc_cb(self,pc_msg):
        
        pc_W_np=pc2_msg_to_np(pc_msg)
        self.pc_W=torch.tensor(pc2_msg_to_np(pc_msg))
        self.pc_cnt_list.append(pc_W_np.shape[1])
        pc_cnt_np=np.array(self.pc_cnt_list)
        np.save(file='pc_cnt1.npy',arr=pc_cnt_np)
        
        self.update_lm()
        
        if self.p.shape[0]!=3:
            return
        
        #init input
        u_init_flatten=self.init_u(view_center='ob_center')
        u_flatten=u_init_flatten
        
        
        optimizer = optim.Adam([u_flatten], lr=0.1)
        #begin optimization
        loss=torch.tensor(0.0)
        
        t=time.time()
        for i in range(20):
            loss_prev=loss
            optimizer.zero_grad()
            loss,vis_pc,vis_ob= step(A=self.A_ext,B=self.B_ext,x0=self.x0,u=u_flatten,p=self.p,pc_W=self.lm_W,n_horizon=self.n_horizon)
            loss.backward()
            
            if torch.isnan(loss):
                print(u_flatten)
                u_flatten=self.init_u(view_center='ob_center')
                
                break
            
            if loss<-1.95 and torch.abs(loss_prev-loss)<1e-3:
                print(i)
                break
            
            optimizer.step()
        
        print(time.time()-t)
        
        
        u_flatten=self.init_u(view_center='ob_center')
        loss,vis_pc,vis_ob= step(A=self.A_ext,B=self.B_ext,x0=self.x0,u=u_flatten,p=self.p,pc_W=self.lm_W,n_horizon=self.n_horizon)
        
        
        
        x_flatten=self.A_ext@self.x0+self.B_ext@u_flatten
        x=x_flatten.reshape([self.n_horizon,2]).T
        
        x_np=x.detach().numpy()
        
        self.x0=torch.tensor(x_np[:,0].reshape([2,1]))
        
        
        ##publish pose array
        euler_zyx=torch.zeros([3,self.n_horizon])
        euler_zyx[:2,:]=x
        T_C_W=euler_position_to_T(euler_zyx=euler_zyx,pos=self.p)
        T_W_C=torch.inverse(T_C_W)
        #(T_W_C.shape)
        #print(torch.rad2deg(self.x0))
        pose_array=PoseArray()
        pose_array.header.frame_id="vrpn"
        for i in range(self.n_horizon):
            T_Bi_V=T_C_W[i,:,:].detach().numpy()
            pose=msgify(Pose,T_Bi_V)
            pose_array.poses.append(pose)

        print(vis_ob)
        print(loss)
        #real time test with pub pointcloud
        self.test_pub.publish(pc_msg)
        self.view_array_pub.publish(pose_array)
        
        
    def pos_arr_cb(self,pos_arr_msg):
        self.p=torch.zeros([3,self.n_horizon])
        
        for i in range(1,self.n_horizon+1):
            self.p[:,i-1]=torch.tensor(numpify(pos_arr_msg.poses[i])[:3,3])
            
        
        
    
    def run(self):
        rate=rospy.Rate(10)
        
        while not rospy.is_shutdown():
            

            rate.sleep()
            
            

if __name__ == '__main__':
    node = ViewPlanningNode()
    node.run()